package com.qisimanxiang.dalaran.zuul.sentinel;

import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowItem;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRuleManager;
import com.qisimanxiang.dalaran.zuul.config.SentinelConfig;
import com.qisimanxiang.dalaran.zuul.util.CommonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.alibaba.csp.sentinel.slots.block.RuleConstant.*;

/**
 * @author wangmeng
 * @date created in 10:12 上午 2019/9/16
 * @modified by
 */
@Slf4j
@Component
public class QpsStrategy {

    private static final int PROVIDER_GLOBAL_DEFAULT_QPS_LIMIT = 1000;
    private static final int PROVIDER_SERVICE_DEFAULT_QPS_LIMIT = 100;
    private static final int PROVIDER_CUSTOMER_SERVICE_DEFAULT_QPS_LIMIT = 10;
    private static final int PROVIDER_CUSTOMER_SERVICE_DEFAULT_TC_LIMIT = 1;
    private static final int MAX_QUEUEING_TIME_MS = 2000;
    private static final int SERVICE_QUEUEING_TIME_MS = 1000;
    private static final int CUSTOMER_QUEUEING_TIME_MS = 200;
    private static final int CUSTOMER_THREAD_QUEUEING_TIME_MS = 2000;
    public static final String RESOURCE_NAME_GLOBAL = "global";
    public static final String RESOURCE_NAME_SERVICE = "service";
    public static final String RESOURCE_NAME_CUSTOMER_SERVICE = "customer-service";

    private static QpsStrategy instance;
    private SentinelConfig sentinelConfig;

    public QpsStrategy() {
        instance = this;
    }

    public static void flush(SentinelConfig config) {
        log.info("refresh qps limit: {}", config);
        instance.sentinelConfig = config;
        instance.flushQpsLimit();
    }

    private void flushQpsLimit() {
        //全局限流配置
        globalQpsLimit();
        //服务级限流配置
        ParamFlowRule serviceFlowRule = serviceQpsLimit(getServiceLimitConfig());
        //客户&服务级限流配置
        ParamFlowRule customerFlowRule = customerServiceQpsLimit(getCustomerServiceLimitConfig());
        //客户线程隔离配置
        ParamFlowRule customerTcFlowRule = customerServiceTcLimit(new HashMap<>());

        //
        List<ParamFlowRule> paramFlowRules = new ArrayList<>();
        paramFlowRules.add(serviceFlowRule);
        paramFlowRules.add(customerFlowRule);
        paramFlowRules.add(customerTcFlowRule);
        ParamFlowRuleManager.loadRules(paramFlowRules);
    }

    private Map<String, Integer> getServiceLimitConfig() {
        return sentinelConfig.getServiceLimit();
    }

    private Map<String, Integer> getCustomerServiceLimitConfig() {
        return sentinelConfig.getCustomerServiceQpsLimit();
    }

    /**
     * 全局限流
     */
    private void globalQpsLimit() {
        FlowRule rule = new FlowRule(RESOURCE_NAME_GLOBAL)
                .setCount(CommonUtils.firstNotNull(sentinelConfig.getGlobalDefaultQpsLimit(), PROVIDER_GLOBAL_DEFAULT_QPS_LIMIT))
                .setGrade(RuleConstant.FLOW_GRADE_QPS)
                //排队并设置最大等待时长
                .setControlBehavior(CONTROL_BEHAVIOR_RATE_LIMITER)
                .setMaxQueueingTimeMs(CommonUtils.firstNotNull(sentinelConfig.getServiceDefaultQueueingMs(), MAX_QUEUEING_TIME_MS));
        List<FlowRule> list = new ArrayList<>();
        list.add(rule);
        FlowRuleManager.loadRules(list);
    }

    /**
     * 服务级别限流
     *
     * @return
     */
    private ParamFlowRule serviceQpsLimit(Map<String, Integer> serviceLimitMap) {
        ParamFlowRule flowRule = new ParamFlowRule(RESOURCE_NAME_SERVICE)
                .setCount(CommonUtils.firstNotNull(sentinelConfig.getServiceDefaultQpsLimit(), PROVIDER_SERVICE_DEFAULT_QPS_LIMIT))
                .setGrade(RuleConstant.FLOW_GRADE_QPS)
                //排队并设置最大等待时长
                .setControlBehavior(CONTROL_BEHAVIOR_RATE_LIMITER)
                .setMaxQueueingTimeMs(CommonUtils.firstNotNull(sentinelConfig.getServiceDefaultQueueingMs(), SERVICE_QUEUEING_TIME_MS))
                .setParamIdx(1);


        List<ParamFlowItem> items = serviceLimitMap.entrySet().stream()
                .map(en -> new ParamFlowItem().setObject(en.getKey()).setClassType(String.class.getName()).setCount(en.getValue()))
                .collect(Collectors.toList());
        flowRule.setParamFlowItemList(items);
        return flowRule;
    }

    /**
     * 客户&服务级别限流
     *
     * @return
     */
    private ParamFlowRule customerServiceQpsLimit(Map<String, Integer> serviceLimitMap) {
        ParamFlowRule flowRule = new ParamFlowRule(RESOURCE_NAME_CUSTOMER_SERVICE)
                .setCount(CommonUtils.firstNotNull(sentinelConfig.getCustomerDefaultQpsLimit(), PROVIDER_CUSTOMER_SERVICE_DEFAULT_QPS_LIMIT))
                .setGrade(RuleConstant.FLOW_GRADE_QPS)
                //策略为直接拒绝
                .setControlBehavior(CONTROL_BEHAVIOR_DEFAULT)
                .setParamIdx(1);

        List<ParamFlowItem> items = serviceLimitMap.entrySet().stream()
                .map(en -> new ParamFlowItem().setObject(en.getKey()).setClassType(String.class.getName()).setCount(en.getValue()))
                .collect(Collectors.toList());
        flowRule.setParamFlowItemList(items);

        return flowRule;
    }



    public static final String RESOURCE_NAME_CUSTOMER_SERVICE_TC = "customer-service-tc";

    /**
     * 客户&服务级别隔离线程
     *
     * @return
     */
    private ParamFlowRule customerServiceTcLimit(Map<String, Integer> customerServiceLimitMap) {
        ParamFlowRule flowRule = new ParamFlowRule(RESOURCE_NAME_CUSTOMER_SERVICE_TC)
                .setCount(CommonUtils.firstNotNull(sentinelConfig.getCustomerDefaultTcLimit(), PROVIDER_CUSTOMER_SERVICE_DEFAULT_TC_LIMIT))
                .setGrade(RuleConstant.FLOW_GRADE_THREAD)
                //策略为等待
                .setControlBehavior(CONTROL_BEHAVIOR_RATE_LIMITER)
                .setMaxQueueingTimeMs(CommonUtils.firstNotNull(sentinelConfig.getCustomerThreadQueueingMs(), CUSTOMER_THREAD_QUEUEING_TIME_MS))
                .setParamIdx(1);

        List<ParamFlowItem> items = customerServiceLimitMap.entrySet().stream()
                .map(en -> new ParamFlowItem().setObject(en.getKey()).setClassType(String.class.getName()).setCount(en.getValue()))
                .collect(Collectors.toList());
        flowRule.setParamFlowItemList(items);
        return flowRule;
    }
}
